#include "../includes/TaskConsumer.h"
#include "../includes/TaskManager.h"
#include <iostream>
#include <thread>
#include <atomic>
#include <chrono>
#include <string>
#include <vector>
#include <sstream>
#include <memory>

int main() {
    // 初始化 TaskManager 和 TaskConsumer
    TaskManager& task_manager = TaskManager::getInstance();
    TaskConsumer& task_consumer = TaskConsumer::getInstance(); // 假设默认构造函数创建10个线程

    std::atomic<bool> running(true); // 用于控制所有线程的运行状态
    std::atomic<int> next_task_id(1); // 用于生成唯一的任务ID

    // --- 线程1: 生产者 - 持续向 TaskManager 添加任务 ---
    std::thread producer_thread([&task_manager, &running, &next_task_id]() {
        int task_counter = 0; // 用于循环发送不同任务
        int num_tasks_to_send = 20; // 总共发送20个任务

        for (int i = 0; i < num_tasks_to_send && running; ++i) {
            std::string task_request;
            int task_id = next_task_id.fetch_add(1);
            int task_type = task_counter % 5; // 循环发送 0, 1, 2, 3, 4 类型的任务

            switch (task_type) {
                case 0: { // 写文件
                    std::string filename = "pool_test_file_" + std::to_string(task_id) + ".txt";
                    std::string content = "Content for pool task " + std::to_string(task_id) + " at iteration " + std::to_string(i) + '\n';
                    
                    if (filename.length() > 128) {
                        std::cerr << "Filename too long for write task." << std::endl;
                        break;
                    }
                    if (content.length() > 1024) {
                        std::cerr << "Content too long for write task." << std::endl;
                        break;
                    }

                    int32_t id_int = task_id;
                    char type_byte = static_cast<char>(task_type);
                    char filename_len_byte = static_cast<char>(filename.length());
                    int32_t content_len_int = content.length();

                    // 构建任务内容 (不包含长度前缀)
                    task_request.append(reinterpret_cast<const char*>(&id_int), 4);
                    task_request.append(&type_byte, 1);
                    task_request.append(&filename_len_byte, 1);
                    task_request.append(filename);
                    task_request.append(reinterpret_cast<const char*>(&content_len_int), 4);
                    task_request.append(content);
                    
                    std::cout << "Producer: Preparing Write Task (ID: " << task_id << ", File: " << filename << ")" << std::endl;
                    break;
                }
                case 1: { // 读文件
                    std::string filename = "pool_test_file_" + std::to_string(task_id - 1) + ".txt"; // 尝试读取上一轮写入的文件
                    
                    if (filename.length() > 128) {
                        std::cerr << "Filename too long for read task." << std::endl;
                        break;
                    }

                    int32_t id_int = task_id;
                    char type_byte = static_cast<char>(task_type);
                    char filename_len_byte = static_cast<char>(filename.length());

                    task_request.append(reinterpret_cast<const char*>(&id_int), 4);
                    task_request.append(&type_byte, 1);
                    task_request.append(&filename_len_byte, 1);
                    task_request.append(filename);
                    
                    std::cout << "Producer: Preparing Read Task (ID: " << task_id << ", File: " << filename << ")" << std::endl;
                    break;
                }
                case 2: { // 创建文件
                    std::string filename = "new_pool_file_" + std::to_string(task_id) + ".txt";
                    
                    if (filename.length() > 128) {
                        std::cerr << "Filename too long for touch task." << std::endl;
                        break;
                    }

                    int32_t id_int = task_id;
                    char type_byte = static_cast<char>(task_type);
                    char filename_len_byte = static_cast<char>(filename.length());

                    task_request.append(reinterpret_cast<const char*>(&id_int), 4);
                    task_request.append(&type_byte, 1);
                    task_request.append(&filename_len_byte, 1);
                    task_request.append(filename);
                    
                    std::cout << "Producer: Preparing Touch Task (ID: " << task_id << ", File: " << filename << ")" << std::endl;
                    break;
                }
                case 3: { // 删除文件
                    std::string filename = "file_to_delete_" + std::to_string(task_id) + ".txt";
                    
                    if (filename.length() > 128) {
                        std::cerr << "Filename too long for rm task." << std::endl;
                        break;
                    }

                    int32_t id_int = task_id;
                    char type_byte = static_cast<char>(task_type);
                    char filename_len_byte = static_cast<char>(filename.length());

                    task_request.append(reinterpret_cast<const char*>(&id_int), 4);
                    task_request.append(&type_byte, 1);
                    task_request.append(&filename_len_byte, 1);
                    task_request.append(filename);
                    
                    std::cout << "Producer: Preparing Rm Task (ID: " << task_id << ", File: " << filename << ")" << std::endl;
                    break;
                }
                case 4: { // 查看目录
                    int32_t id_int = task_id;
                    char type_byte = static_cast<char>(task_type);

                    task_request.append(reinterpret_cast<const char*>(&id_int), 4);
                    task_request.append(&type_byte, 1);
                    
                    std::cout << "Producer: Preparing Ls Task (ID: " << task_id << ")" << std::endl;
                    break;
                }
                default:
                    break;
            }

            if (!task_request.empty()) {
                task_manager.add_task(task_request);
                std::cout << "Producer: Added task (ID: " << task_id << ") to TaskManager queue." << std::endl;
            }

            task_counter++;
            // 模拟任务产生频率
            std::this_thread::sleep_for(std::chrono::milliseconds(200));
        }
        std::cout << "Producer thread finished." << std::endl;
    });

    // --- 线程2: 分发者 - 持续从 TaskManager 取任务并提交给 TaskConsumer ---
    std::thread dispatcher_thread([&task_manager, &task_consumer, &running]() {
        while (running) {
            // 检查 TaskConsumer 是否有空闲线程
            if (task_consumer.idle_threads() > 0 && !task_manager.task_empty()) {
                std::string task_str;
                // 将任务提交给 TaskConsumer 执行
                if (!task_consumer.add_task()) { // 注意：这里调用的是无参数版本，它会从 TaskManager 获取任务
                        std::cerr << "Dispatcher: Failed to add task to TaskConsumer." << std::endl;
                } else {
                    std::cout << "Dispatcher: Submitted task to TaskConsumer." << std::endl;
                }
            } else {
                // 没有空闲线程，短暂休眠
                std::this_thread::sleep_for(std::chrono::milliseconds(10));
            }
        }
        std::cout << "Dispatcher thread finished." << std::endl;
    });

    // --- 线程3: 消费者 - 持续从 TaskManager 读取响应并打印 ---
    std::thread consumer_thread([&task_manager, &running]() {
        std::string accumulated_data; // 用于处理TCP粘包 (如果响应也是长度前缀的话)

        while (running) {
            if (!task_manager.resp_empty()) { // 检查 TaskManager 响应队列是否非空
                std::string resp_str;
                task_manager.get_resp(resp_str); // 假设 get_resp 是阻塞的或非阻塞但配合 sleep

                if (!resp_str.empty()) {
                    // --- 解析响应内容 (假设内容是 [任务ID(4字节) + 类型ID(1字节) + 状态码(1字节) + ...]) ---
                    if (resp_str.length() < 6) { // 至少需要 4 + 1 + 1
                        std::cerr << "Consumer: Received response content is too short." << std::endl;
                        continue; // 尝试继续接收下一个响应
                    }

                    size_t offset = 0;
                    int task_id = *(reinterpret_cast<const int*>(resp_str.data() + offset));
                    offset += 4;
                    char type_id = resp_str[offset];
                    offset += 1;
                    char status_code = resp_str[offset];
                    offset += 1;

                    std::cout << "Consumer: Received Response - Task ID: " << task_id << ", Type: " << static_cast<int>(type_id) 
                              << ", Status Code: " << static_cast<int>(status_code);

                    std::string response_details = "";

                    switch (type_id) {
                        case 0: // 写文件响应: +4字节bytes_written
                            if (resp_str.length() >= offset + 4) {
                                int bytes_written = *(reinterpret_cast<const int*>(resp_str.data() + offset));
                                response_details = ", Bytes Written: " + std::to_string(bytes_written);
                            } else {
                                std::cerr << " [ERROR] Write response content is truncated." << std::endl;
                            }
                            break;
                        case 1: // 读文件响应: +4字节content_length + 变长content
                            if (resp_str.length() >= offset + 4) {
                                int32_t content_len = *(reinterpret_cast<const int*>(resp_str.data() + offset));
                                if (resp_str.length() >= offset + 4 + content_len) {
                                    std::string content_str = resp_str.substr(offset + 4, content_len);
                                    response_details = ", Content Length: " + std::to_string(content_len) + ", Content: " + content_str;
                                } else {
                                    std::cerr << " [ERROR] Cat response content is truncated." << std::endl;
                                }
                            } else {
                                std::cerr << " [ERROR] Cat response content is truncated." << std::endl;
                            }
                            break;
                        case 2: // 创建文件响应: 无额外内容
                        case 3: // 删除文件响应: 无额外内容
                            // No extra details to print beyond status code
                            break;
                        case 4: // 查看目录响应: +1字节file_count + 变长file_entries
                            if (resp_str.length() >= offset + 1) {
                                char file_count = resp_str[offset];
                                offset += 1;
                                std::string entries_str = resp_str.substr(offset);
                                std::istringstream iss(entries_str);
                                std::string file_name;
                                std::vector<std::string> files;
                                while (std::getline(iss, file_name)) {
                                    files.push_back(file_name);
                                }
                                response_details = ", File Count: " + std::to_string(static_cast<int>(file_count)) + ", Files: [";
                                for(size_t i = 0; i < files.size(); ++i) {
                                    response_details += files[i];
                                    if (i < files.size() - 1) response_details += ", ";
                                }
                                response_details += "]";
                            } else {
                                std::cerr << " [ERROR] Ls response content is truncated." << std::endl;
                            }
                            break;
                        default:
                            std::cerr << " [ERROR] Unknown response type ID: " << static_cast<int>(type_id) << std::endl;
                            // 尝试继续，但可能数据已损坏
                            break;
                    }

                    std::cout << response_details << std::endl;
                }
            } else {
                // 响应队列为空，短暂休眠避免忙等待
                std::this_thread::sleep_for(std::chrono::milliseconds(10));
            }
        }
        std::cout << "Consumer thread finished." << std::endl;
    });

    // 等待生产者线程完成（发送完所有任务）
    producer_thread.join();
    std::cout << "All tasks sent by producer. Waiting for processing..." << std::endl;

    // 给一些时间让剩余的任务被处理完
    std::this_thread::sleep_for(std::chrono::seconds(2));

    // 停止其他线程
    running = false;

    // 等待所有线程结束
    if (dispatcher_thread.joinable()) {
        dispatcher_thread.join();
    }
    if (consumer_thread.joinable()) {
        consumer_thread.join();
    }

    std::cout << "All threads finished. Test completed." << std::endl;
    return 0;
}